package zookeeper

import (
	"context"
	"encoding/json"
	"errors"
	"gitee.com/vrv_media/go-micro-framework/registry"
	regOps "gitee.com/vrv_media/go-micro-framework/registry/options"
	"github.com/go-zookeeper/zk"
	"golang.org/x/sync/singleflight"
	"path"
	"time"
)

// Option is etcd registry option.
type Option func(o *options)

type options struct {
	rootPath string
	user     string
	password string
}

// ZookeeperRegistry is consul registry
type ZookeeperRegistry struct {
	opts *options
	conn *zk.Conn

	group singleflight.Group
}

func NewDefaultZookeeperRegistry(reg *regOps.RegistryOptions, opts ...Option) *ZookeeperRegistry {
	conn, _, err := zk.Connect([]string{reg.Address}, time.Second*15)
	if err != nil {
		panic(err)
	}
	return NewZookeeperRegistry(conn, opts...)
}

func NewZookeeperRegistry(conn *zk.Conn, opts ...Option) *ZookeeperRegistry {
	options := &options{
		rootPath: "/microservices",
	}
	if len(opts) > 0 {
		for _, o := range opts {
			o(options)
		}
	}

	return &ZookeeperRegistry{
		opts: options,
		conn: conn,
	}
}

func (r *ZookeeperRegistry) Register(_ context.Context, service *registry.ServiceInstance) error {
	var (
		data []byte
		err  error
	)
	if err = r.ensureName(r.opts.rootPath, []byte(""), 0); err != nil {
		return err
	}
	serviceNamePath := path.Join(r.opts.rootPath, service.Name)
	if err = r.ensureName(serviceNamePath, []byte(""), 0); err != nil {
		return err
	}
	if data, err = marshal(service); err != nil {
		return err
	}
	servicePath := path.Join(serviceNamePath, service.ID)
	if err = r.ensureName(servicePath, data, zk.FlagEphemeral); err != nil {
		return err
	}
	go r.reRegister(servicePath, data)
	return nil
}

// Deregister registry service to zookeeper.
func (r *ZookeeperRegistry) Deregister(ctx context.Context, service *registry.ServiceInstance) error {
	ch := make(chan error, 1)
	servicePath := path.Join(r.opts.rootPath, service.Name, service.ID)
	go func() {
		err := r.conn.Delete(servicePath, -1)
		ch <- err
	}()
	var err error
	select {
	case <-ctx.Done():
		err = ctx.Err()
	case err = <-ch:
	}
	return err
}

// reRegister re-register data node info when bad connection recovered
func (r *ZookeeperRegistry) reRegister(path string, data []byte) {
	sessionID := r.conn.SessionID()
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()
	for range ticker.C {
		cur := r.conn.SessionID()
		// sessionID changed
		if cur > 0 && sessionID != cur {
			// re-ensureName
			if err := r.ensureName(path, data, zk.FlagEphemeral); err != nil {
				return
			}
			sessionID = cur
		}
	}
}

// ensureName ensure node exists, if not exist, create and set data
func (r *ZookeeperRegistry) ensureName(path string, data []byte, flags int32) error {
	exists, stat, err := r.conn.Exists(path)
	if err != nil {
		return err
	}
	// ephemeral nodes handling after restart
	// fixes a race condition if the server crashes without using CreateProtectedEphemeralSequential()
	if flags&zk.FlagEphemeral == zk.FlagEphemeral {
		err = r.conn.Delete(path, stat.Version)
		if err != nil && !errors.Is(err, zk.ErrNoNode) {
			return err
		}
		exists = false
	}
	if !exists {
		if len(r.opts.user) > 0 && len(r.opts.password) > 0 {
			_, err = r.conn.Create(path, data, flags, zk.DigestACL(zk.PermAll, r.opts.user, r.opts.password))
		} else {
			_, err = r.conn.Create(path, data, flags, zk.WorldACL(zk.PermAll))
		}
		if err != nil {
			return err
		}
	}
	return nil
}
func marshal(si *registry.ServiceInstance) ([]byte, error) {
	return json.Marshal(si)
}

func unmarshal(data []byte) (si *registry.ServiceInstance, err error) {
	err = json.Unmarshal(data, &si)
	return
}

// WithRootPath with registry root path.
func WithRootPath(path string) Option {
	return func(o *options) {
		o.rootPath = path
	}
}

// WithDigestACL with registry password.
func WithDigestACL(user string, password string) Option {
	return func(o *options) {
		o.user = user
		o.password = password
	}
}
